from flask import Flask, request, jsonify, Response,stream_with_context
from config import Config
from utils import problem_enhancements
from model import db, init_db, VisitCount
import json
import time
import os
import shutil

token = "eyJ0eXBlIjoiSldUIiwiYWxnIjoiSFM1MTIifQ.eyJqdGkiOiI6IlJPTEVfUkVHSVNURVIiLCJpc3MiOiJPcGVuWExhYiIsImlhdCI6MTczNDM0NjQ4NSwiY2xpZW50SWQiOiJlYm1ydm9kNnlvMG5semFlazF5cCIsInBob25lIjoiMTg1MDg0MzY4OTMiLCJ1dWlkIjoiMTQ3MjkxNDUtNDcyZC00NGQzLWI4NWItNzgzNmY2OGQxMTY2IiwiZW1haWwiOiIiLCJleHAiOjE3NDk4OTg0ODV9.QX4S4jqXk3aI3_my_b5BEfTg02Rm8yLyYqabm8m4y5MYTDRO1lsf2_7hcABcXhMwoAbDC3izSvWIf5vZzwFvSQ"


app = Flask(__name__)
app.config.from_object(Config)
db.init_app(app)

# Create tables within application context
with app.app_context():
    db.create_all()

# --------------------------- 测试----------------------
@app.route('/api/')
def index():
    return 'Hello World!'

# --------------------------- 访问量----------------------
@app.route('/api/visit_count', methods=['GET'])
def visit_count():
    visit_count = VisitCount.query.first()
    
    if not visit_count:
        visit_count = VisitCount(count=1)
        db.session.add(visit_count)
    else:
        visit_count.count += 1
    
    db.session.commit()
    return jsonify({'visit_count': visit_count.count})


# --------------------------- 问题增强----------------------

@app.route('/api/problem_enhancement', methods=['POST'])
def problem_enhancement():
    data = request.get_json()
    user_input = data['user_input']
    if 'token' not in data:
        token = token
    else:
        token = data['token']
    # 增强问题判断是不是只有一端
    n = 0
    while True:
        data = problem_enhancements(user_input, token)
        if len(data.split('\n')) == 1:
            break
        else:
            n += 1
            if n > 3:
                break
   
    return jsonify({'searchQuery': data.split('\n')[0].strip()})


# --------------------------- 语义检索----------------------
@app.route('/api/search')
def search():
    query = request.args.get('query', '')

    import uuid
    import os
    path = f"data/{uuid.uuid4()}/milvus_demo.db"
    if os.path.exists("data"):
        os.makedirs(os.path.dirname(path), exist_ok=True)
    else:
        os.makedirs(os.path.dirname(path))
    def generate():
        steps = [
            "分析查询意图...",
            "生成检索词...",
            "执行文献检索...",
            "计算相关性得分...",
            "结果排序与筛选...",
            "生成摘要与推荐..."
        ]
        # from utils import 
        # for i, step in enumerate(steps[:-1]):
        #     yield f"data: {json.dumps({'step': i, 'message': step}, ensure_ascii=False)}\n\n"
        #     time.sleep(1)  # 模拟处理时间
        # ----------------------1 分析意图----------------------
        i= 0
        from utils import intent_recognition
        intent = intent_recognition(query,token)
        # # 如果True在intent中，则返回结果
        # if 'True' in intent:
        yield f"data: {json.dumps({'step': i, 'message': steps[i]}, ensure_ascii=False)}\n\n"
        time.sleep(1)
        # ----------------------2 生成检索词----------------------
        i += 1
        from utils import generate_search_keywords
        search_keywords = generate_search_keywords(question = query,token=token)
        yield f"data: {json.dumps({'step': i, 'message': steps[i]}, ensure_ascii=False)}\n\n"
        time.sleep(2)
        # ----------------------3 执行检索----------------------
        i += 1
        from utils import execute_search
        search_results = execute_search(question=query,token=token,search_keys=search_keywords) 
        yield f"data: {json.dumps({'step': i, 'message': steps[i]}, ensure_ascii=False)}\n\n"
        time.sleep(2)

        # ----------------------4 计算相关性得分----------------------
        i += 1
        from utils import cal_embedding
        cal_embedding(results=search_results,path=path)
        yield f"data: {json.dumps({'step': i, 'message': steps[i]}, ensure_ascii=False)}\n\n"
        time.sleep(2)
        # ----------------------5 结果排序与筛选----------------------
        i += 1
        from utils import rank_and_filter
        results,index_list = rank_and_filter(question=query,token=token,path=path,top_n=20)
        yield f"data: {json.dumps({'step': i, 'message': steps[i]}, ensure_ascii=False)}\n\n"
        time.sleep(2)

        # ----------------------6 生成报告----------------------
        i += 1
        from utils import generate_report
        report = generate_report(question=query,token=token,results=results,index_list=index_list,top_n=10)
        yield f"data: {json.dumps({'step': i, 'message': steps[i]}, ensure_ascii=False)}\n\n"
        time.sleep(2)

        # 最后一步，返回结果
        results = {
            "report": report,
            "results": results
            }
        try:
            if os.path.exists(os.path.dirname(path)):
                shutil.rmtree(os.path.dirname(path), ignore_errors=True)
        except Exception:
            pass
        # 最终结果作为一个完整的响应返回
        yield f"data: {json.dumps({'step': 5, 'message': steps[-1], 'results': results, 'done': True}, ensure_ascii=False)}\n\n"
    


    return Response(generate(), mimetype='text/event-stream')


if __name__ == '__main__':
    app.run(use_reloader=False)